from shutil import copy2
# 自己刚刚转换出来的txt路径
datadir_normal = "./fire_smoke/all_label" 
# 所有的图像路径
datadir_normal1 = './fire_smoke/JPEGImages'
all_data = os.listdir(datadir_normal)  # 图片文件夹
num_all_data = len(all_data)
print("总数量: " + str(num_all_data))
index_list = list(range(num_all_data))
print(index_list)

# 生成随机种子,使得每一次划分的数据一样
r = random.random
random.seed(2)
random.shuffle(index_list, random=r)

num = 0
# 新建两个目录,分别是labels和images
trainDir = "./fire_smoke/labels/train"
# 将训练集txt放在这个文件夹下
if not os.path.exists(trainDir):
    os.mkdir(trainDir)
validDir = './fire_smoke/labels/val 
# 将测试集txt放在这个文件夹下
if not os.path.exists(validDir):
    os.mkdir(validDir)
trainDir1 = './fire_smoke/images/train'
# 将训练集图片放在这个文件夹下
if not os.path.exists(trainDir1):
    os.mkdir(trainDir1)
validDir1 = './fire_smoke/images/val 
# 将验证集图片放在这个文件夹下
if not os.path.exists(validDir1):
    os.mkdir(validDir1)
for i in index_list:
    fileName = os.path.join(datadir_normal, all_data[i])
    if num < num_all_data * 0.8:  # Train:0.8
        copy2(fileName, trainDir)
    else:  # val:0.2
        copy2(fileName, validDir)
    num += 1
num = 0
for i in index_list:
    fileName = os.path.join(datadir_normal1, all_data[i].replace("txt","jpg"))
    if num < num_all_data * 0.8:
        copy2(fileName, trainDir1)
    else:  # val:0.2
        copy2(fileName, validDir1)
    num += 1 
